import cv2
import numpy as np
from skimage.segmentation import flood


def mask_flow_with_det_boxes(flow, det_boxes):
    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
    mag = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
    mag = mag.astype(dtype=np.uint8, copy=True)

    flow_h, flow_w, _ = flow.shape
    for det_box in det_boxes:
        x1, y1, x2, y2 = det_box
        c_x, c_y = (x1 + x2) / 2, (y1 + y2) / 2
        c_x = int(min(max(0, c_x * flow_w), flow_w - 1))
        c_y = int(min(max(0, c_y * flow_h), flow_h - 1))
        mask = flood(mag, (c_y, c_x), tolerance=5)
        flow[mask] = 0
    return flow


class DenseOpticalFlow(object):
    def __init__(self, step=16):
        self.step = step

    def optical_flow(self, prev_img, curr_img, trk_boxes=None):
        flow = cv2.calcOpticalFlowFarneback(prev=prev_img,
                                            next=curr_img,
                                            flow=None,
                                            pyr_scale=0.5,
                                            levels=5,
                                            winsize=11,
                                            iterations=5,
                                            poly_n=5,
                                            poly_sigma=1.1,
                                            flags=0)
        if trk_boxes is not None:
            flow = mask_flow_with_det_boxes(flow, trk_boxes)
        # 使用最小二乘法计算Homography
        flow_h, flow_w, flow_c = flow.shape
        y, x = np.mgrid[self.step:flow_h - self.step:self.step,
                        self.step:flow_w - self.step:self.step].reshape(2, -1)
        fx, fy = flow[y, x].T
        flag = np.ones_like(y)[:, None]
        A = np.hstack([x[:, None], y[:, None], flag])
        y = np.hstack([(x + fx)[:, None], (y + fy)[:, None], flag])
        H = np.matmul(np.linalg.inv(np.matmul(A.T, A)), np.matmul(A.T, y))
        H = H.T
        return H
